////////////////////////////////////////////////////////////////////////////////////////////////////
// file:	C:\Users\Quan\Documents\Visual Studio 2010\Projects\source\cvWavelet\cvWavelet2D.cpp
//
// summary:	Implements the cv wavelet 2 d class
////////////////////////////////////////////////////////////////////////////////////////////////////

#include <string> 
#include <cufft.h> 
#include <cuda.h> 
#include <iostream> 
#include <fstream>
#include <iomanip>
#include <opencv2/opencv.hpp>
#include "dwt.hpp"
#include "cutils.hpp"


using namespace std;
using namespace cv;

////////////////////////////////////////////////////////////////////////////////////////////////////
/// <summary>	Main entry-point for this application. </summary>
///
/// <remarks>	Quan, 5/9/2012. </remarks>
///
/// <param name="argc">	Number of command-line arguments. </param>
/// <param name="argv">	Array of command-line argument strings. </param>
///
/// <returns>	Exit-code for the process - 0 for success, else an error code. </returns>
////////////////////////////////////////////////////////////////////////////////////////////////////
int main(int argc, char* argv[])
{
	//cudaSafeCall(cudaSetDevice(2));
	//cudaSafeCall(cudaSetDevice(4));
	cudaSafeCall(cudaDeviceReset());
	
	/// <summary>	Declare the image input file. </summary>
	const char* iFile;

	/// <summary>	By default choose Lenna image in the solution directory. </summary>
	if(argc==1)
		iFile = "lenna.jpg";
	else 
		iFile = argv[1];

	/// <summary>	The image. </summary>
	Mat img;

	/// <summary>	Read image in grayscale mode. </summary>
	img = imread(iFile, IMREAD_GRAYSCALE);

	/// <summary>	If image does not exist, terminate program. </summary>
	if (img.empty())
	{
		cout << "Cannot read image" << endl;
		return -1;
	}

	/// <summary>	Display the source image. </summary>
	//imshow("Input", img); 

	/// <summary>	Retrieve the number of rows and columns. </summary>
	int nRows = img.rows;
	int nCols = img.cols;
	
	/// <summary>	Allocate 1d array of data. </summary>
	float2* data;
	data = (float2*)malloc(nRows*nCols*sizeof(float2));

	/// <summary>	Copy data from image to 1d array of data. </summary>
	for(int i=0; i<nRows; i++)
	{
		for(int j=0; j<nCols; j++)
		{
			data[i*nCols + j].x = (float)img.at<uchar>(i, j);
			data[i*nCols + j].y = 0;
		}
	}
	

	cudaEvent_t beginFWT, beginIWT;
	cudaEvent_t endFWT,   endIWT;
	float 		timeFWT,	timeIWT;
	float 		tFWT,		tIWT;

	timeFWT  	= 0.0f;
	timeIWT  	= 0.0f;

		
	cudaSafeCall(cudaEventCreate(&beginFWT));
	cudaSafeCall(cudaEventCreate(&beginIWT));


	
	cudaSafeCall(cudaEventCreate(&endFWT));
	cudaSafeCall(cudaEventCreate(&endIWT));

	
	/// <summary>	Copy data from CPU to GPU </summary>

	float2 *d_data;
	
	cudaSafeCall(
		cudaMalloc((void**)&d_data, sizeof(float2)*nCols*nRows)
		);
	cudaSafeCall(
		cudaMemcpy2D(d_data, sizeof(float2)*nCols,  
			data, sizeof(float2)*nCols, 
			sizeof(float2)*nCols, nRows, 
			cudaMemcpyHostToDevice)
		);
	
	////////////////////////////////////////////////////////////////////////////////////////////////////
	/// <summary>	Initial plan for wavelet.		</summary>
	float2 *tmp_transpose;
	
	cudaMalloc((void**)&tmp_transpose,	sizeof(float2)*nRows*nCols);

	////////////////////////////////////////////////////////////////////////////////////////////////////
	/// <summary>	Perform Forward Wavelet Transform </summary>
	float2 *d_wavelet;
	cudaSafeCall(
		cudaMalloc((void**)&d_wavelet, sizeof(float2)*nCols*nRows)
		);
	////////////////////////////////////////////////////////////////////////////////////////////////////
	cudaSafeCall(cudaEventRecord(beginFWT, 0));
	cuda::dwt(d_data, 
		d_wavelet, 
		nRows, nCols, 
		cuda::DWT_FORWARD,
		tmp_transpose);
	cudaSafeCall(cudaEventRecord(endFWT, 0));
	cudaSafeCall(cudaEventSynchronize(endFWT));
	cudaSafeCall(cudaEventElapsedTime(&tFWT, beginFWT, endFWT));
	timeFWT += tFWT;
	////////////////////////////////////////////////////////////////////////////////////////////////////	
	float2* wavelet;
	wavelet = (float2*)malloc(nRows*nCols*sizeof(float2));
	cudaSafeCall(
		cudaMemcpy2D(wavelet, sizeof(float2)*nCols,  
		d_wavelet, sizeof(float2)*nCols, 
		sizeof(float2)*nCols, nRows, 
		cudaMemcpyDeviceToHost)
		);
	Mat waveletImg(nRows, nCols, CV_32FC1);
	for(int i=0; i<nRows; i++)
	{
		for(int j=0; j<nCols; j++)
		{
			waveletImg.at<float>(i, j) = wavelet[i*nCols + j].x;
		}
	}
	waveletImg.convertTo(waveletImg, CV_8UC1);
	//imshow("Forward Wavelet", waveletImg);
	imwrite("Forward Wavelet.tif", waveletImg);
	/// <summary>	Perform Inverse Wavelet Transform </summary>
	float2 *d_inverse;
	cudaSafeCall(
		cudaMalloc((void**)&d_inverse, sizeof(float2)*nCols*nRows)
		);
	
	////////////////////////////////////////////////////////////////////////////////////////////////////
	cudaSafeCall(cudaEventRecord(beginIWT, 0));
	
	cuda::dwt(d_wavelet, 
		d_inverse, 
		nRows, nCols, 
		cuda::DWT_INVERSE,
		tmp_transpose);
	cudaSafeCall(cudaEventRecord(endIWT, 0));
	cudaSafeCall(cudaEventSynchronize(endIWT));
	cudaSafeCall(cudaEventElapsedTime(&tIWT, beginIWT, endIWT));
	timeIWT += tIWT;
	////////////////////////////////////////////////////////////////////////////////////////////////////
	float2* inverse;
	inverse = (float2*)malloc(nRows*nCols*sizeof(float2));
	cudaSafeCall(
		cudaMemcpy2D(inverse, sizeof(float2)*nCols,  
		d_inverse, sizeof(float2)*nCols, 
		sizeof(float2)*nCols, nRows, 
		cudaMemcpyDeviceToHost)
		);
		
	Mat inverseImg(nRows, nCols, CV_32FC1);
	for(int i=0; i<nRows; i++)
	{
		for(int j=0; j<nCols; j++)
		{
			inverseImg.at<float>(i, j) = inverse[i*nCols + j].x;
		}
	}
	inverseImg.convertTo(inverseImg, CV_8UC1);
	//imshow("Inverse Wavelet", inverseImg);
	imwrite("Inverse Wavelet.tif", inverseImg);
	
	cout << "FWT:    \t" << timeFWT 	<< "ms" << endl;
	cout << "IWT:    \t" << timeIWT 	<< "ms" << endl;
	return 0;
}